﻿import time
from FlagEmbedding import FlagReranker
# reranker = FlagReranker('E:/study/bge-m3/bge-large-zh-v1.5', use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation

# score = reranker.compute_score(['query', 'passage'])
# print(score) # -5.65234375

# # You can map the scores into 0-1 by set "normalize=True", which will apply sigmoid function to the score
# score = reranker.compute_score(['query', 'passage'], normalize=True)
# print(score) # 0.003497010252573502

# scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
# print(scores) # [-8.1875, 5.26171875]
# T1 = time.time()
# datas=[];
# datas.append(['本院认为，被告人余勇无视国法，以非法占有为目的，入户盗窃，其行为触犯了《中华人民共和国刑法》第二百六十四条之规定，犯罪事实清楚，证据确实、充分，应当以盗窃罪追究其刑事责任', '夜间盗窃']);
# datas.append(['本院认为，被告人余勇无视国法，以非法占有为目的，入户盗窃，其行为触犯了《中华人民共和国刑法》第二百六十四条之规定，犯罪事实清楚，证据确实、充分，应当以盗窃罪追究其刑事责任', '入户盗窃']);
# datas.append(['绍兴市柯桥区人民检察院指控，2022年7月12日凌晨，被告人盛国庆到绍兴市柯桥区柯桥街道双川小区15幢楼下，采用撬门方式进入该幢15号的仓库，窃得蒋来放在该处的中华香烟、利群香烟、冬虫夏草香烟等共计540余条。经价格认定，涉案香烟共计价值219721元。', '入户盗窃']);
# datas.append(['绍兴市柯桥区人民检察院指控，2022年7月12日凌晨，被告人盛国庆到绍兴市柯桥区柯桥街道双川小区15幢楼下，采用撬门方式进入该幢15号的仓库，窃得蒋来放在该处的中华香烟、利群香烟、冬虫夏草香烟等共计540余条。经价格认定，涉案香烟共计价值219721元。', '夜间盗窃']);
# # You can map the scores into 0-1 by set "normalize=True", which will apply sigmoid function to the score
# scores = reranker.compute_score(datas, normalize=True)
# print(scores) # [0.00027803096387751553, 0.9948403768236574]
# T2 = time.time()
# print('程序运行时间:%s毫秒' % ((T2 - T1)*1000))


from FlagEmbedding import FlagModel
sentences_1 = ["本院认为，被告人杜某某醉酒后在道路上无证驾驶机动车，被追尾造成他人死亡，其行为构成危险驾驶罪。"]
sentences_2 = ["数额特别巨大","入户盗窃","多次盗窃","在医院盗窃病人亲友财物","曾因盗窃受过刑事处罚又实施盗窃","携带凶器盗窃","结伙盗窃","盗窃毒品","盗窃家庭成员的财物","盗窃违禁品","控制未成年人盗窃","扒窃","其他严重情节","夜间盗窃","数额巨大","为吸毒而盗窃","在公共交通工具上盗窃","数额较大","造成公私财产损失","携带管制刀具盗窃","在公共场所盗窃","一年内曾因盗窃受过行政处罚","盗窃信用卡并使用","采用破坏性手段盗窃","无情形情节","流窜盗窃","在医院盗窃病人财物","其他特别严重情节"]
model = FlagModel('E:/study/bge-m3/bge-large-zh-v1.5', 
                  query_instruction_for_retrieval="盗窃罪情形判断",
                  use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
embeddings_2 = model.encode(sentences_2)
T1 = time.time()
embeddings_1 = model.encode(sentences_1)
T2 = time.time()
print('程序运行时间:%s毫秒' % ((T2 - T1)*1000))
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

# for s2p(short query to long passage) retrieval task, suggest to use encode_queries() which will automatically add the instruction to each query
# corpus in retrieval task can still use encode() or encode_corpus(), since they don't need instruction
queries = ["本院认为，被告人杜某某醉酒后在道路上无证驾驶机动车，被追尾造成他人死亡，其行为构成危险驾驶罪。"]
passages = ["被查处时驾车逃跑","被查处时抗拒检查","被查处时让人顶替","被查处时逃跑","曾因酒后驾驶受过刑事追究又危险驾驶","曾因酒后驾驶受过行政处罚又危险驾驶","从事旅客运输危险驾驶","从事校车业务危险驾驶","负事故全部责任","负事故同等责任","负事故主要责任","驾驶大型普通客车","驾驶客车","驾驶摩托车","驾驶使用变造车牌证的机动车","驾驶使用伪造车牌证的机动车","驾驶危险品运输车","驾驶无号牌的机动车","驾驶小型普通客车","驾驶已达报废标准的机动车","驾驶营运车","驾驶中型普通客车","三年内曾因酒后驾驶受过刑事追究","事故后逃逸","逃避公安机关依法检查","通用情形","违反交通信号灯通行","违反危险化学品安全管理规定运输危险化学品危及公共安全","无驾驶资格","无情形情节","严重超速行驶","严重超员行驶","严重超载行驶","饮酒驾驶","在城市快速路上驾驶","在高速公路上驾驶","造成交通事故","致人轻伤","致人轻微伤","致人死亡","追逐竞驶情节恶劣","醉酒驾驶"]
q_embeddings = model.encode_queries(queries)
p_embeddings = model.encode(passages)
scores = q_embeddings @ p_embeddings.T
print(scores)